import numpy as np
import acado
import time
import matplotlib.pyplot as plt
from numpy import sin, cos,exp


def weighted_center_of_pc(p_body,pc_W):
    pc_B=pc_W-p_body
    dist_2=np.sum(pc_B*pc_B,axis=0)
    dist=np.sqrt(dist_2)

    weight=(1/dist)/np.sum(1/dist)
    pc_W_weighted=weight*pc_W
    pc_W_weighted_center=np.sum(pc_W_weighted,axis=1)
    print(pc_W_weighted_center)

    return pc_W_weighted_center

pc_W=np.load("../ov_pc_W_left_bottom corner.npy")

print(pc_W.shape)




#test    
p_body=np.array([1,2,3]).reshape([3,1])
pc_W=np.ones([3,100])*2


#weighted_center_of_pc(p_body,pc_W)
    


pc_size=4
NPC=4

N=20 #timestep
NU=3
NOD=4+3*NPC
NX=6
NY=7
NYN=4
PC_X_START=4
PC_Y_START=PC_X_START+NPC
PC_Z_START=PC_X_START+2*NPC

x0=np.zeros([1,NX])
x0[0,1]=np.deg2rad(45)
x0[0,2]=np.arctan(1/10)
X=np.tile(x0.flatten(),(N+1,1))
x0=np.zeros([1,NX])

#x0=np.array([100,2,3,1,0,0,0,0,0,0]).reshape(1,NX)
#X=np.zeros([N+1,NX])+0.
u=np.zeros([N,NU])+0.0

OD=np.ones([N+1,NOD])

Q=np.eye(NY)
Q[0,0]=1e9 #v1
Q[1,1]=1e0 #v5
Q[2,2]=1e0 #v6
Q[3,3]=1e0 #a1
Q[4,4]=1e0 #a5
Q[5,5]=1e0 #a6
Q[6,6]=1e10 #cost
W=np.transpose(np.tile(Q,N))
WN=np.eye(NYN)*1
WN[0,0]=Q[0,0]
WN[1,1]=Q[1,1]
WN[2,2]=Q[2,2]
WN[3,3]=Q[6,6]



y=np.zeros([N,NY])
#y[:,:3]=yN[:,:3]

OD_i=np.array([10.0001,10.0001,0.001,0,10.001,10])

a=0.666

OD=np.ones([N+1,NOD])*1e5


###
OD[:,0]=a #a
OD[0,1:4]=[0,0,0] #bx,by,bz (body)

OD[1:,1:4]=[0,0,0] #cmd traj

pc_W_x=10.1
OD[:,PC_X_START:PC_X_START+pc_size]=pc_W_x+np.random.randn(pc_size)*0.0 #pc_x

pc_W_y=pc_W_x*0.1
OD[:,PC_Y_START:PC_Y_START+pc_size]=pc_W_y+np.random.randn(pc_size)*0.0#pc_y

pc_W_z=pc_W_x
OD[:,PC_Z_START:PC_Z_START+pc_size]=pc_W_z+np.random.randn(pc_size)*0.0#pc_z

t0=time.time()
X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)

x0=X[-1,:].reshape([1,6])
#X=np.zeros([N+1,NX])+0.
#u=np.zeros([N,NU])+0.0
#X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)


plot=1

if plot:

    plt.subplot(221)
    plt.plot(np.rad2deg(X[:,:3]))
    plt.legend(['p1','p5','p6'])
    plt.grid()
    plt.subplot(222)
    plt.plot(X[:,3:6])
    plt.legend(['v1','v5','v6'])
    plt.grid()
    plt.subplot(223)

    plt.plot(u)
    plt.legend(['a1','a5','a6'])
    plt.grid()

    plt.subplot(224)
    
    
    bear=np.arange(-100,100,0.01)
    bear_rad=np.deg2rad(bear)
    v=1/(np.exp(-a*np.cos(bear_rad)+3)+1);

    cos_bear=np.cos(bear_rad)
    k2=-0.24;
    k1= 0.44;
    k0= 0.69;

    v=k0+k1*cos_bear+k2*cos_bear*cos_bear
    
    plt.plot(bear,np.abs(bear_rad))
    plt.plot(bear,-np.log(v))

    #plt.plot(OD[0,:])
    #plt.ylim([-20,20])
    plt.grid()
    plt.show()


